"""流式响应智能体"""

import asyncio
from typing import AsyncGenerator, List, Dict
from json_repair import repair_json
from loguru import logger
import traceback

from model.excutor_model import ExecutionStatus
from agent.react_agent import ReactAgent
from agent.agent_state import AgentState
from schemas.interaction import *


class AdjustedAction(Action):
    input: Dict = Field(default={}, description='动作输入参数')

class AdjustedActionBegin(ActionBegin):
    data: AdjustedAction = Field(description='action的详细数据')

class AdjustedActionEnd(ActionEnd):
    data: AdjustedAction = Field(description='action的详细数据')


class StreamingReactAgent(ReactAgent):
    """实时响应的流式智能体"""
    
    def __init__(self, verbose: bool = False, mode: str = "hybrid", enable_reflection: bool = True):
        super().__init__(verbose, mode, enable_reflection=enable_reflection)
        self._event_queues = dict()
        self._current_steps = dict()
        self._current_plan_ids = dict()
        self._current_plans = dict()
        self._current_actions = dict()
    
    async def run_stream(
            self, session_id: str = None,
            query: str = None,
            interrupt_feedback: str = None,
            history: List[Any] = None,
            max_steps: int = None
    ) -> AsyncGenerator[InteractionBase, None]:

        self._event_queues[session_id] = asyncio.Queue()

        # 创建计划
        current_plan_id = f"plan_{uuid.uuid4()}"
        self._current_plan_ids[session_id] = current_plan_id
        self._current_steps[session_id] = 0

        current_plan = Plan(
            id=current_plan_id,
            sort_index=self._current_steps[session_id],
            name="处理用户提问",
            status="pending",
            actions=[],
        )
        self._current_plans[session_id] = current_plan

        await self._emit_event(AddPlan(
            session_id=session_id,
            data=[current_plan]
        ))

        # 在后台任务中启动智能体执行
        task = asyncio.create_task(
            self._run_with_events(session_id, query, interrupt_feedback, history, max_steps)
        )
        try:
            while True:
                try:
                    event = await asyncio.wait_for(self._event_queues[session_id].get(), timeout=1.0)
                    logger.info(event)
                    await asyncio.sleep(0)
                    yield event
                    if event.name == 'Finished':
                        break
                except asyncio.TimeoutError:
                    if task.done():
                        break
                    continue
        except Exception as e:
            logger.error(traceback.format_exc())
    
    async def _run_with_events(
            self,
            session_id: str = None,
            query: str = None,
            interrupt_feedback: str = None,
            history: List[Any] = None,
            max_steps: int = None
    ):
        try:
            await self.run(session_id, query, interrupt_feedback, history, max_steps)
        except Exception as e:
            logger.error(traceback.format_exc())
            await self._emit_event(Finished(
                session_id=session_id,
                status='failed',
                result=str(e),
            ))

    async def _emit_event(self, event: InteractionBase):
        if self._event_queues.get(event.session_id, None):
            await self._event_queues[event.session_id].put(event)

    async def _decide_approach_node(self, state: AgentState) -> AgentState:

        session_id = state.get('session_id')
        current_plan_id = self._current_plan_ids[session_id]
        # 添加动作
        self._current_steps[session_id] += 1
        current_action = AdjustedAction(
            flowResult=[],
            id=f"action_{uuid.uuid4()}",
            name="选择处理方式",
            input={"query": state.get('input')},
            plan_id=current_plan_id,
            sort_index=self._current_steps[session_id],
            status='pending'
        )

        # 执行动作
        await self._emit_event(AdjustedActionBegin(
            session_id=state.get('session_id'),
            id = current_action.id,
            data = current_action
        ))

        result_state = await super()._decide_approach_node(state)

        current_action.flowResult = [
            {
                'method': result_state['chosen_approach']
            }
        ]
        current_action.status = 'completed'

        await self._emit_event(AdjustedActionEnd(
            session_id=state.get('session_id'),
            id=current_action.id,
            data=current_action
        ))

        return result_state

    async def _think_node(self, state: AgentState) -> AgentState:

        session_id = state.get('session_id')
        result_state = await super()._think_node(state)

        # 若某项行动已规划，则发送行动已规划事件
        if result_state["actions"] and len(result_state["actions"]) > len(result_state["observations"]):
            latest_action = result_state["actions"][-1]

            name = latest_action.get("name", latest_action.get("action", "unknown"))

            action_input = latest_action.get("input", latest_action.get("action_input", {}))
            if not isinstance(action_input, dict):
                try:
                    action_input = repair_json(str(action_input), return_objects=True)
                except Exception as e:
                    logger.info(e)
            if not isinstance(action_input, dict):
                action_input = {}

            self._current_steps[session_id] = state['current_step'] + 1
            current_action = AdjustedAction(
                flowResult = [],
                id = f"action_{uuid.uuid4()}",
                name = f"调用工具: {name}",
                input = action_input,
                plan_id = self._current_plan_ids[session_id],
                sort_index=self._current_steps[session_id],
                status = 'pending'
            )
            self._current_actions[session_id] = current_action

        return result_state
    
    async def _act_node(self, state: AgentState) -> AgentState:
        if not state["actions"]:
            return state

        session_id = state.get('session_id')
        if session_id in self._current_actions:
            current_action = self._current_actions[session_id]
            await self._emit_event(AdjustedActionBegin(
                session_id=session_id,
                id = current_action.id,
                data = current_action
            ))

        result_state = await super()._act_node(state)

        # 发送动作执行结果
        if result_state["tool_results"]:
            latest_result = result_state["tool_results"][-1]
            tool_result = latest_result.get("result", {})

            current_action = self._current_actions[session_id]
            current_action.flowResult = [tool_result]

            if isinstance(tool_result, dict):
                current_action.status = 'completed' if tool_result.get('success', False) else 'failed'
            else:
                current_action.status = 'completed'

            await self._emit_event(AdjustedActionEnd(
                session_id=state.get('session_id'),
                id=current_action.id,
                data=current_action
            ))

        return result_state

    async def _plan_node(self, state: AgentState) -> AgentState:

        result_state = await super()._plan_node(state)

        if result_state.get("current_plan"):
            plan = result_state["current_plan"]

            session_id = state.get('session_id')
            current_plan_id = self._current_plan_ids[session_id]

            actions = []
            for i, step in enumerate(getattr(plan, 'steps', [])):
                self._current_steps[session_id] += 1
                action = AdjustedAction(
                    flowResult=[],
                    id=step.id,
                    name=f"{step.id}: {getattr(step, 'description', '无')}",
                    input={},
                    sort_index=self._current_steps[session_id],
                    plan_id=current_plan_id,
                    status='pending',
                )
                actions.append(action)

            self._current_plans[session_id].actions = actions

        return result_state
    
    async def _execute_node(self, state: AgentState) -> AgentState:
        """计划执行有多种模式：并行/串行/迭代等"""

        id2action = {}
        session_id = state.get('session_id')
        for action in self._current_plans[session_id].actions:
            await self._emit_event(AdjustedActionBegin(
                session_id=session_id,
                id=action.id,
                data=action
            ))
            id2action[action.id] = action

        result_state = await super()._execute_node(state)

        if state.get("execution_result"):
            for i, step_result in enumerate(state["execution_result"].step_results):
                step_out = step_result.output
                if not isinstance(step_out, dict):
                    step_out = {"output": str(step_out)}

                action = id2action[step_result.step_id]
                action.flowResult = [step_out]
                if step_result.status == ExecutionStatus.COMPLETED:
                    action.status = "completed"
                else:
                    action.status = "failed"
                id2action[step_result.step_id] = action

        for action_id, action in id2action.items():
            if action.status not in ["completed", "failed"]:
                action.status = "failed"

            await self._emit_event(AdjustedActionEnd(
                session_id=session_id,
                id=action.id,
                data=action
            ))
            await asyncio.sleep(0)

        return result_state
    
    async def _reflect_node(self, state: AgentState) -> AgentState:
        if not self.reflection_module:
            return state
        session_id = state.get('session_id')
        current_plan_id = self._current_plan_ids[session_id]
        self._current_steps[session_id] += 1
        current_action = AdjustedAction(
            flowResult = [],
            id = f"action_{uuid.uuid4()}",
            name = f"过程反思",
            input = {},
            plan_id = current_plan_id,
            sort_index=self._current_steps[session_id],
            status = 'pending'
        )
        self._current_actions[session_id] = current_action

        await self._emit_event(AdjustedActionBegin(
            session_id=state.get('session_id'),
            id=current_action.id,
            data=current_action
        ))

        result_state = await super()._reflect_node(state)
        reflection_metadata = result_state.get("metadata", {}).get("reflection", {})

        current_action.flowResult.append({
            "reply_score": f"{reflection_metadata.get('final_quality_score', 0.0):.1f}"
        })
        current_action.status = 'completed'

        await self._emit_event(AdjustedActionEnd(
            session_id=state.get('session_id'),
            id=current_action.id,
            data=current_action
        ))

        return result_state

    async def _finish_node(self, state: AgentState) -> AgentState:

        final_state = await super()._finish_node(state)
        output = final_state.get("output", "") if isinstance(final_state, dict) else ""
        success = not final_state.get("has_error", False) if isinstance(final_state, dict) else False
        session_id = state.get('session_id')
        await self._emit_event(Finished(
            session_id=session_id,
            status='completed' if success else 'failed',
            result=output
        ))
        logger.debug(f"最终回答:\n{output}")
        return final_state